#include "kernel_operator.h"
#include "op_common.h"
namespace AscendC {

template <typename _DT_X1, typename _DT_X2, typename _DT_Y> class Kernel {
public:
  TPipe pipe;
  int64_t size = 0;
  int64_t tileLength = 0;
  int64_t finalLength = 0;
  DefInTensor(X1);
  DefInTensor(X2);
  DefOutTensor(Y);

//   DefBufVECCALC(CASTED_X);
//   DefBufVECCALC(CASTED_Y);

  DefBufVECCALC(TEMP);
//   DefBufVECCALC(TEMP2);

public:
  __aicore__ inline Kernel() {}

  template <typename T>
  __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, T tiling_data) {
    this->size = tiling_data.size;
    this->tileLength = tiling_data.tileLength;
    this->finalLength = tiling_data.finalLength;
    // 切分global
    GSetBuffer(X1, x1, 0, tiling_data.size);
    GSetBuffer(X2, x2, 0, tiling_data.size);
    GSetBuffer(Y, y, 0, tiling_data.size);
    // 初始化队列
    InitQueueSimple(X1, tileLength);
    InitQueueSimple(X2, tileLength);
    InitQueueSimple(Y, tileLength);
    // 初始化BUF
    InitTBufBuffer(TEMP, tileLength * sizeof(_DT_X1));
    TBufGet(TEMP, _DT_X1);
    Duplicate(BTensor(TEMP), _DT_X1(0), tileLength);
    // InitTBufBuffer(TEMP2, tileLength * sizeof(float));
    // if constexpr (std::is_same_v<_DT_X, half>) {
    //   InitTBufBuffer(CASTED_X, tileLength * sizeof(float));
    //   InitTBufBuffer(CASTED_Y, tileLength * sizeof(float));
    // }
  }

  __aicore__ inline void Process() {
    auto loopCount = size / tileLength;
    for (uint32_t i = 0; i < loopCount; ++i) {
      CopyIn(i, tileLength);
      Compute(tileLength);
      CopyOut(i, tileLength);
    }
    if (finalLength > 0) {
      auto copyCount = alignToBlock<_DT_X1>(finalLength);
      CopyIn(loopCount, copyCount);
      Compute(finalLength);
      CopyOut(loopCount, copyCount);
    }
  }

  __aicore__ inline void CopyIn(uint32_t i, uint32_t calcCount) {
    EnQueGlobal2Local(X1, i * tileLength, calcCount);
    EnQueGlobal2Local(X2, i * tileLength, calcCount);
  }

  __aicore__ inline void CopyOut(uint32_t i, uint32_t calcCount) {
    DeQueLocal2Global(Y, i * tileLength, calcCount);
  }

  __aicore__ inline void Compute(uint32_t calcCount) {
    // if constexpr (std::is_same_v<_DT_X, half>) {      
    // } else {      
    // }
    TBufGet(TEMP, _DT_X1);
    DeQueSimple(X1);
    DeQueSimple(X2);
    QueAllocSimple(Y);
    Ln(LTensor(X2), LTensor(X2), calcCount);
    Mul(LTensor(Y), LTensor(X1), LTensor(X2), calcCount);
    auto x2_int8 = LTensor(X2).template ReinterpretCast<int8_t>();
    Compare(x2_int8, LTensor(X1), BTensor(TEMP), CMPMODE::NE, calcCount);
    Select(LTensor(Y), x2_int8, LTensor(Y), _DT_X1(0), SELMODE::VSEL_TENSOR_SCALAR_MODE, calcCount);
    EnQue(Y);
    QueFree(X1);
    QueFree(X2);
  }
};
} // namespace AscendC
extern "C" __global__ __aicore__ void xlogy(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
     AscendC::Kernel<DTYPE_X1, DTYPE_X2, DTYPE_Y> op;
    op.Init(x1, x2, y, tiling_data);
    op.Process();
}